/*
 * Copyright 2015 MingxingWang

 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at

 *     http://www.apache.org/licenses/LICENSE-2.0

 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.audaque.graphdb.index

import java.io._

import com.audaque.graphdb._
import org.apache.lucene.analysis.Analyzer
import org.apache.lucene.analysis.standard.StandardAnalyzer
import org.apache.lucene.document._
import org.apache.lucene.index.IndexWriterConfig.OpenMode
import org.apache.lucene.index._
import org.apache.lucene.search.BooleanClause.Occur
import org.apache.lucene.search._
import org.apache.lucene.store._
import org.apache.lucene.util.Version
import org.apache.spark.graphx.VertexId

import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer

/**
 * 索引类，提供索引有关的各种索引接口，只能为一个图的顶点创建索引，不提供边的索引
 * 这里的索引是由Apache Lucene技术支持，其原理请参见Apache Lucene官网
 * Created by mingxingwang on 2015/3/14.
 */
class LuceneIndex(val indexPath: File) {
  val indexDataFile = "index"
  val indexInfoFile = "_META.txt"
  val fieldNames: mutable.Map[Long, Set[String]] = new mutable.HashMap[Long, Set[String]]()
  var schemaFactory: SchemaFactory = null
  var openWay = "RAM"

  var dfsDir: Directory = null
  var analyzer: Analyzer = null
  var iwc: IndexWriterConfig = null
  var indexWriter: IndexWriter = null
  var indexSearcher: IndexSearcher = null
  private val idFieldName = "_ID_"
  private val fuzzyThreshold = 0.1

  /**
   * 初始化索引操作
   * @param fields  指定在哪些域上创建索引，同时初始化IndexWriter，以备后面进行索引写操作
   */
  def init(schFact: SchemaFactory, fields: Map[Long, Set[String]]): Unit = {
    this.fieldNames.clear()
    if (fields == null || fields.isEmpty) {
      return
    }
    this.fieldNames ++= fields
    this.schemaFactory = schFact
    dfsDir = FSDirectory.open(new File(indexPath, indexDataFile))
    analyzer = new StandardAnalyzer()
    iwc = new IndexWriterConfig(Version.LUCENE_4_10_1, analyzer)
    iwc.setOpenMode(OpenMode.CREATE_OR_APPEND)
    indexWriter = new IndexWriter(dfsDir, iwc)
  }

  /**
   * 将一个顶点Vertex的信息添加到索引中
   * @param vid  表示一个顶点Vertex在图数据库中的查找码
   * @param prop  包含一个Vertex的属性名称和属性值
   * @return
   */
  def add(vid: VertexId, prop: Property): Boolean = {
    if (prop == null || prop.schema == null) return true
    val sch = prop.schema
    val doc: Document = new Document
    var bIndexed = false
    fieldNames.get(sch.id) match {
      case Some(names) =>
        names.foreach(fieldName => {
          val t = sch.getFieldType(fieldName)
          if (t != Type.UnKnown) {
            val value = prop(fieldName)
            if (value != null) {
              add(doc, value, fieldName, t)
              bIndexed = true
            }
          }
        })
      case None => return false
    }
    if (bIndexed) {
      val idField: LongField = new LongField(idFieldName, vid, Field.Store.YES)
      doc.add(idField)
      indexWriter.addDocument(doc)
    }
    true
  }

  def add(doc: Document, fieldValue: Any, fieldName: String, fieldType: Type): Unit = {
    if (fieldValue == null) return
    val t = LuceneIndex.TypeId2Type.getOrElse(fieldType.tid, 0)
    if (t == 1) {
      val field: DoubleField = new DoubleField(fieldName.toLowerCase, fieldValue.toString.toDouble, Field.Store.NO)
      doc.add(field)
    } else if (t == 2) {
      val field: LongField = new LongField(fieldName.toLowerCase, fieldValue.toString.toLong, Field.Store.NO)
      doc.add(field)
    } else {
      val field: StringField = new StringField(fieldName.toLowerCase, fieldValue.toString.toLowerCase, Field.Store.NO)
      doc.add(field)
    }
  }

  /**
   * 索引是否处于打开状态，索引打开才能基于索引查询
   * @return
   */
  def isOpen: Boolean = {
    indexSearcher != null
  }

  /**
   * 打开索引
   * @param openWay  打开索引的方式，“ROM”加载到内存和“RAM”从磁盘打开两种方式
   */
  def openIndex(schFact: SchemaFactory, openWay: String = "RAM"): Unit = {
    if (indexSearcher != null)
      close()
    this.schemaFactory = schFact
    getInfo()
    if (openWay.equalsIgnoreCase("ROM"))
      openIndexSearcher_ROM()
    else if (openWay.equalsIgnoreCase("RAM"))
      openIndexSearcher_RAM()
    else
      throw new IllegalArgumentException(String.format("Open Index Way Must Be: \"RAM\" or \"ROM\"."))
  }

  /**
   * 打开索引查询并加载到内存，打开一次，使用多次
   */
  private def openIndexSearcher_RAM(): Unit = {
    if (dfsDir == null) {
      dfsDir = FSDirectory.open(new File(indexPath, indexDataFile))
    }
    if (indexSearcher != null)
      indexSearcher.getIndexReader.close()
    val ramDir: RAMDirectory = new RAMDirectory(dfsDir, new IOContext())
    val indexReader: IndexReader = DirectoryReader.open(ramDir)
    indexSearcher = new IndexSearcher(indexReader)
    openWay = "RAM"
  }

  /**
   * 直接从磁盘打开索引，打开一次，使用多次
   */
  private def openIndexSearcher_ROM(): Unit = {
    if (dfsDir == null) {
      dfsDir = FSDirectory.open(new File(indexPath, indexDataFile))
    }
    if (indexSearcher != null)
      indexSearcher.getIndexReader.close()
    val indexReader: IndexReader = DirectoryReader.open(dfsDir)
    indexSearcher = new IndexSearcher(indexReader)
    openWay = "ROM"
  }

  def getInfo() {
    val metaFile = new File(indexPath, indexInfoFile)
    if (!metaFile.exists()) {
      return
    }
    fieldNames.clear()
    val in: DataInputStream = new DataInputStream(new FileInputStream(metaFile))
    val size = in.readInt()
    for (i <- 0 until size) {
      val sid = in.readLong()
      val ns = in.readInt()
      val names = new mutable.HashSet[String]()
      for (j <- 0 until ns) {
        names += in.readUTF()
      }
      fieldNames += (sid -> names.toSet)
    }
  }

  /**
   * 全字段查询多个关键字
   * @param values  待查询的关键字集合
   * @param logic  logic为真各关键字按照“与”规则查询，否则按照“或”规则查询
   * @param fuzzyLookup  为真进行模糊查询，否则精确查询
   * @return
   */
  def valueSearch(values: Iterable[String], logic: Boolean = false, fuzzyLookup: Boolean = false, retNum: Int = 100): Array[VertexId] = {
    val id_list: ArrayBuffer[VertexId] = new ArrayBuffer[VertexId]()
    val booleanQuery: BooleanQuery = new BooleanQuery()
    var occur: Occur = null
    if (logic)
      occur = Occur.MUST
    else
      occur = Occur.SHOULD
    for (value <- values) {
      var query: Query = null
      if (fuzzyLookup)
        query = getGlobalFuzzyQuery(value)
      else
        query = getGlobalPreciseQuery(value)
      if (query != null) {
        booleanQuery.add(query, occur)
      }
    }
    val docs: TopDocs = indexSearcher.search(booleanQuery, retNum)
    val scoreDocs: Array[ScoreDoc] = docs.scoreDocs
    for (doc <- scoreDocs) {
      val document: Document = indexSearcher.doc(doc.doc)
      id_list += document.get(idFieldName).toLong
    }
    id_list.toArray
  }

  /**
   * 键值对的查找方法
   * @param namesAndValues  需要查询的键值对的集合
   * @param logic            需要查询的键值对按“与”还是按“或”查询
   * @param fuzzyLookup   是否按照模糊查询，为true表示模糊查询，否则为精确查询
   * @param retNum           最多返回多少条查询结果
   * @return
   */
  def pairSearch(namesAndValues: Iterable[(String, String)], logic: Boolean = false, fuzzyLookup: Boolean = false, retNum: Int = 100): Array[VertexId] = {
    val id_list: ArrayBuffer[VertexId] = new ArrayBuffer[VertexId]
    val booleanQuery: BooleanQuery = new BooleanQuery()
    var occur: Occur = null
    if (logic)
      occur = Occur.MUST
    else
      occur = Occur.SHOULD
    for ((name, value) <- namesAndValues) {
      if (name != null && value != null) {
        for ((sid, names) <- fieldNames) {
          schemaFactory.getSchemaById(sid) match {
            case Some(sch) =>
              if (names.contains(name)) {
                val fieldType = sch.getFieldType(name)
                val t = LuceneIndex.TypeId2Type.getOrElse(fieldType.tid, 0)
                val query = if (fuzzyLookup) {
                  genFuzzyQuery(name, value, t)
                } else {
                  genPreciseQuery(name, value, t)
                }
                if (query != null) {
                  booleanQuery.add(query, occur)
                }
              }
            case None =>
          }
        }
      }
    }
    val docs: TopDocs = indexSearcher.search(booleanQuery, retNum)
    val scoreDocs: Array[ScoreDoc] = docs.scoreDocs
    for (doc <- scoreDocs) {
      val document: Document = indexSearcher.doc(doc.doc)
      id_list += document.get(idFieldName).toLong
    }
    id_list.toArray
  }

  private def genFuzzyQuery(name: String, value: String, t: Int): Query = {
    if (value == null) {
      return null
    }
    try {
      if (t == 1) {
        val doubleValue = value.toDouble
        val min = doubleValue * (1 - fuzzyThreshold)
        val max = doubleValue * (1 + fuzzyThreshold)
        NumericRangeQuery.newDoubleRange(name.toLowerCase, min, max, true, true)
      } else if (t == 2) {
        val longValue = value.toLong
        val min = (longValue * (1 - fuzzyThreshold)).toLong
        val max = (longValue * (1 + fuzzyThreshold)).toLong
        NumericRangeQuery.newLongRange(name.toLowerCase, min, max, true, true)
      } else {
        new FuzzyQuery(new Term(name.toLowerCase, value.toLowerCase))
      }
    } catch {
      case ex: NumberFormatException => null
    }
  }

  private def genPreciseQuery(name: String, value: String, t: Int): Query = {
    if (name == null || value == null) {
      return null
    }
    try {
      if (t == 1) {
        val doubleValue = value.toDouble
        NumericRangeQuery.newDoubleRange(name.toLowerCase, doubleValue - 1e-6, doubleValue + 1e-6, true, true)
      } else if (t == 2) {
        val longValue = value.toLong
        NumericRangeQuery.newLongRange(name.toLowerCase, longValue, longValue, true, true)
      } else {
        new TermQuery(new Term(name.toLowerCase, value.toLowerCase))
      }
    } catch {
      case ex: NumberFormatException => null
    }
  }

  /**
   * 根据一个查询值返回和该值对应的全字段查询(模糊查询)
   * @param value  查询值
   * @return  和查询值对应的全字段查询(模糊查询)
   */
  def getGlobalFuzzyQuery(value: String): Query = {
    if (value == null) {
      return null
    }
    val booleanQuery: BooleanQuery = new BooleanQuery()
    for ((sid, names) <- fieldNames) {
      schemaFactory.getSchemaById(sid) match {
        case Some(sch) =>
          for (name <- names) {
            val fieldType = sch.getFieldType(name)
            val t = LuceneIndex.TypeId2Type.getOrElse(fieldType.tid, 0)
            val query = genFuzzyQuery(name, value, t)
            if (query != null) {
              booleanQuery.add(query, Occur.SHOULD)
            }
          }
        case None =>
      }
    }
    booleanQuery
  }

  /**
   * 根据一个查询值返回和该值对应的全字段查询(精确查询)
   * @param value  查询值
   * @return  和查询值对应的全字段查询(精确查询)
   */
  def getGlobalPreciseQuery(value: String): Query = {
    if (value == null) {
      return null
    }
    val booleanQuery: BooleanQuery = new BooleanQuery()
    for ((sid, names) <- fieldNames) {
      schemaFactory.getSchemaById(sid) match {
        case Some(sch) =>
          for (name <- names) {
            val fieldType = sch.getFieldType(name)
            val t = LuceneIndex.TypeId2Type.getOrElse(fieldType.tid, 0)
            val query = genPreciseQuery(name, value, t)
            if (query != null) {
              booleanQuery.add(query, Occur.SHOULD)
            }
          }
        case None =>
      }
    }
    booleanQuery
  }

  def commit(): Unit = {
    if (indexWriter == null) {
      return
    }
    val metaFile = new File(indexPath, indexInfoFile)
    indexPath.mkdirs()
    val out: DataOutputStream = new DataOutputStream(new FileOutputStream(metaFile))
    out.writeInt(fieldNames.size)
    for ((sid, names) <- fieldNames) {
      out.writeLong(sid)
      out.writeInt(names.size)
      for (name <- names) {
        out.writeUTF(name)
      }
    }
    out.flush()
    out.close()
    indexWriter.close()
  }


  /**
   * 关闭索引
   */
  def close(): Unit = {
    if (indexSearcher != null) {
      indexSearcher.getIndexReader.close()
      indexSearcher = null
    }
  }
}

object LuceneIndex {
  val TypeId2Type: Map[Int, Int] = Map(Type.ID_Double -> 1,
    Type.ID_Float -> 1,
    Type.ID_Int -> 2,
    Type.ID_Long -> 2,
    Type.ID_Short -> 2)
}